import numpy as np
import matplotlib.pyplot as plt


from read_data import *
from write_data import *
from plot_data import *
from globals import *

from plot_data import calc_min_f

MODEL = "F22"
RESULT_DIR = "../OUTPUT_Filter/"
file_list = ["A33"]


def plot1(input, angle, freq_motion,freq_angle):
    t_start = 0
    t_end = len(input)

    # 画一个周期
    f_min = calc_min_f(freq_motion)
    print("最小频率："+str(f_min)+"Hz")
    t_start = int(len(input)/2 - 1/f_min * freq_angle/2)
    t_end = int(len(input)/2 + 1/f_min * freq_angle/2)
    input = input[t_start:t_end,:]
    angle = angle[t_start:t_end,:]

    print(t_start/freq_angle,t_end/freq_angle)

    alpha = angle[:,0]
    beta = angle[:,1]

    headers = ["Fx", "Fy", "Fz", "Mx", "My", "Mz", "Q"]

    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(alpha,input[:,i])
            plt.xlabel('Alpha')
            plt.ylabel('')
    else:
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(beta,input[:,i])
            plt.xlabel('Beta')
            plt.ylabel('')


def comp_result(file_list):
    F_n_Q_aero={}
    F_n_Q_body={}
    F_n_Q_wind={}
    F_n_Q_rslt={}
    C_n_Q_body={}
    C_n_Q_wind={}
    C_n_Q_rslt={}
    Angle_body={}
    Angle_wind={}
    

    for item in file_list:
        PREFIX_DIR = RESULT_DIR + MODEL+ "/" + item

        F_n_Q_aero[item] = read_d_result(PREFIX_DIR + "/force_aero.csv")
        F_n_Q_body[item] = read_d_result(PREFIX_DIR + "/force_body.csv")
        F_n_Q_wind[item] = read_d_result(PREFIX_DIR + "/force_wind.csv")
        F_n_Q_rslt[item] = read_d_result(PREFIX_DIR + "/force_rslt.csv")

        C_n_Q_body[item] = read_d_result(PREFIX_DIR + "/coeff_body.csv")
        C_n_Q_wind[item] = read_d_result(PREFIX_DIR + "/coeff_wind.csv")
        C_n_Q_rslt[item] = read_d_result(PREFIX_DIR + "/coeff_rslt.csv")

        Angle_body[item] = read_d_result(PREFIX_DIR + "/angle_body.csv")
        Angle_wind[item] = read_d_result(PREFIX_DIR + "/angle_wind.csv")


    # 画图
    
    plt.figure(figsize=FIG_SIZE)  # Set the size of the figure
    t = np.linspace(0, 20, 2000)

    # Plotting wind angles
    for i in range(2):  # Loop through the first two types of wind angles
        for item in file_list:  # Iterate through each file in the list
            plt.subplot(2, 2, i + 1)  # Create a subplot for each wind angle type
            plt.plot(t[0:len(Angle_body[item][:, i])],Angle_wind[item][:, i], label=('Pitch/Roll' if item == file_list[0] else 'Pitch/Roll/Yaw'))  # Plot the wind angle, first line is 'a', second is 'b'
            plt.xlabel('Time (s)')  # Label for the x-axis
            plt.ylabel(['Angle of Attack (degrees)', 'Sideslip Angle (degrees)'][i])  # Label for the y-axis
            plt.legend()  # Add a legend to the subplot

    # Plotting body angles
    for i in range(3):  # Loop through the three types of body angles
        for item in file_list:  # Iterate through each file in the list
            plt.subplot(2, 3, i + 1 + 3)  # Create a subplot for each body angle type
            plt.plot(t[0:len(Angle_body[item][:, i])],Angle_body[item][:, i], label=('Pitch/Roll' if item == file_list[0] else 'Pitch/Roll/Yaw'))  # Plot the body angle, first line is 'a', second is 'b'
            plt.xlabel('Time (s)')  # Label for the x-axis
            plt.ylabel(['Yaw (degrees)', 'Pitch (degrees)', 'Roll (degrees)'][i])  # Label for the y-axis
            plt.legend()  # Add a legend to the subplot

    plt.show()  # Display the plots 
        
#     # 角度对比
#     plt.figure(figsize=FIG_SIZE)
#     for i in range(2):
#         for item in file_list:
#             plt.subplot(2,2,i+1)
#             plt.plot(Angle_wind[item][:,i])
#     for i in range(3):
#         for item in file_list:
#             plt.subplot(2,3,i+1+3)
#             plt.plot(Angle_body[item][:,i])
#     plt.show()
    

#     # 3维轨迹
#     # 创建一个图形实例
#     cc = ["Pitch+Roll", "Pitch+Roll+Yaw"]
#     fig = plt.figure()
#     # 在图形中创建一个3D子图
#     ax = fig.add_subplot(111, projection='3d')

#     # 遍历文件列表，绘制每个文件的三维曲线
#     for i, item in enumerate(file_list):
#         pitch_w = Angle_body[item][:,1]
#         roll_w = Angle_body[item][:,0]
#         yaw_w = Angle_body[item][:,2]

#         # 绘制曲线，并为每条曲线指定一个图例标签
#         ax.plot(roll_w, yaw_w, pitch_w, label=cc[i])

#     # 设置坐标轴标签和图例
#     ax.set_xlabel('yaw')
#     ax.set_ylabel('roll')
#     ax.set_zlabel('pitch')
#     ax.legend()  # 确保图例显示

#     # 调整布局并显示图形
#     plt.tight_layout()
#     plt.show()

    # # 气动力对比
    # plt.figure(figsize=FIG_SIZE)
    # for i in range(6):
    #     for item in file_list:
    #         plt.subplot(2,3,i+1)
    #         plt.plot(F_n_Q_body[item][:,i])
    # plt.show()

    # plt.figure(figsize=FIG_SIZE)
    # for i in range(6):
    #     for item in file_list:
    #         plt.subplot(2,3,i+1)
    #         plt.plot(F_n_Q_wind[item][:,i])
    # plt.show()

#     # 角度
    plt.figure(figsize=FIG_SIZE)
    for i in range(6):
        for item in file_list:
            plt.subplot(2,3,i+1)
            plt.plot(Angle_wind[item][:,0], C_n_Q_rslt[item][:,i])
    plt.show()

#     # plt.figure(figsize=FIG_SIZE)
#     # for i in range(6):
#     #     for item in file_list:
#     #         plt.subplot(2,3,i+1)
#     #         plt.plot(Angle_wind[item][:,0], F_n_Q_body[item][:,i])
#     # plt.show()
#     DATASET_PREFIX = "../DATASET_ALL/DATASET_UNCUT/"
#     freq_motion_dict = get_f(MODEL, DATASET_PREFIX)
#     plt.figure(figsize=FIG_SIZE)
#     for item in file_list:
#         freq_motion = freq_motion_dict[item]
#         plot1(C_n_Q_rslt[item],Angle_wind[item],freq_motion,100)
#     plt.show()


comp_result(file_list)